#include <stdlib.h>
#include <time.h>
#include <stdio.h>

#define OK 0
#define ERR 1

enum COLOR{RED,BLACK};
typedef struct RBTreeNode
{
	int key;
	enum COLOR color;
	struct RBTreeNode *lchild;
	struct RBTreeNode *rchild;
	struct RBTreeNode *p;
	
}RBTreeNode;

typedef struct RBTree
{
	struct RBTreeNode *pnil;
	struct RBTreeNode *proot;
}RBTree;


void tracePrint(RBTreeNode *p,RBTreeNode *flag,int size){
	RBTreeNode **arr = (RBTreeNode**)malloc(sizeof(RBTreeNode*)*size);
	
	if(p!=flag){
		arr[0] = p;
	}
	int len = 1;
	RBTreeNode *temp;
	int i=0;

	while(len>0){
		int slen = len;
		while(slen>0){
			temp = arr[0];
			for(i=0;i<len-1;i++){
				arr[i]=arr[i+1];
			}
			arr[len--]=NULL;
			printf("%d(color=%s) ",temp->key,temp->color==RED?"RED":"BLACK");
			if(temp->lchild!=flag){
				arr[len] = temp->lchild;
				len++;
			}
			if(temp->rchild!=flag){
				arr[len] = temp->rchild;
				len++;
			}
			slen--;
		}
		printf("\n");
	}
	free(arr);
}

void leftRotate(RBTree *T,RBTreeNode *x)
{
	RBTreeNode *y = x->rchild;
	x->rchild = y->lchild;
	if(x->lchild!=T->pnil){
		y->lchild->p=x;
	}
	y->p = x->p;
	if(x->p==T->pnil){
		T->proot = y;
	}else if(x==x->p->lchild){
		x->p->lchild = y;
	}else{
		x->p->rchild = y;
	}
	y->lchild = x;
	x->p = y;
}

void rightRotate(RBTree *T,RBTreeNode *x)
{
	RBTreeNode *y = x->lchild;
	x->lchild = y->rchild;
	if(y->rchild!=T->pnil){
		y->rchild->p=x;
	}
	y->p = x->p;
	if(x->p==T->pnil){
		T->proot = y;
	}else if(x==x->p->lchild){
		x->p->lchild = y;
	}else{
		x->p->rchild = y;
	}
	y->rchild = x;
	x->p = y;
}

void fixupRBTree(RBTree *T,RBTreeNode *z){

	RBTreeNode *y;
	while(z->p->color==RED){
		
		if(z->p==z->p->p->lchild){
			y = z->p->p->rchild;
			if(y->color == RED){
				z->p->color = BLACK;
				y->color = BLACK;
				z->p->p->color = RED;
				z = z->p->p;
			}else{
				if(z==z->p->rchild){
					z = z->p;
					leftRotate(T,z);
				}
				z->p->color = BLACK;
				z->p->p->color = RED;
				rightRotate(T,z->p->p);
			}
		}else{
			y = z->p->p->lchild;
			if(y->color == RED){
				z->p->color = BLACK;
				y->color = BLACK;
				z->p->p->color = RED;
				z = z->p->p;
			}else{
				if(z==z->p->lchild){
					z = z->p;
					rightRotate(T,z);
				}
				z->p->color = BLACK;
				z->p->p->color = RED;
				leftRotate(T,z->p->p);
			}
		}
	}
	T->proot->color = BLACK;
	
}

int insertRBTree(RBTree *T,int key){

	RBTreeNode *y = T->pnil;
	RBTreeNode *x = T->proot;
	
	while(x!=T->pnil){
		y = x;
		if(x->key>key){
			x = x->lchild;
		}else if(x->key<key){
			x = x->rchild;
		}else{
			return ERR;
		}
	}

	RBTreeNode *z = (RBTreeNode*)malloc(sizeof(RBTreeNode));
	z->key = key;
	z->color = RED;
	z->p = y;
	z->lchild = T->pnil;
	z->rchild = T->pnil;
	
	if(y==T->pnil){
		T->proot = z;
	}else if(y->key>z->key){
		y->lchild = z;
	}else{
		y->rchild = z;
	}
	fixupRBTree(T,z);
	return OK;
}

void rbTransplant(RBTree *tree,RBTreeNode *u,RBTreeNode *v){
	if(u->p==tree->pnil){
		tree->proot = v;
	}else if(u==u->p->lchild){
		u->p->lchild = v;
	}else{
		u->p->rchild = v;
	}
	
	v->p = u->p;
}
void rbDeleteFixup(RBTree *tree,RBTreeNode *x){
	RBTreeNode *w;
	while(x!=tree->proot&&x->color==BLACK){
		if(x==x->p->lchild){
			w = x->p->rchild;
			if(w->color==RED){
				w->color = BLACK;
				x->p->color = RED;
				leftRotate(tree,x->p);
				w = x->p->rchild;
			}
			if(w->lchild->color==BLACK&&w->rchild->color==BLACK){
				w->color = RED;
				x = x->p;
			}else{
				if(w->rchild->color==BLACK){
					w->lchild->color = BLACK;
					w->color = RED;
					rightRotate(tree,w);
					w = x->p->rchild;
				}
				w->color = x->p->color;
				x->p->color = BLACK;
				w->rchild->color = BLACK;
				leftRotate(tree,x->p);
				x = tree->proot;
			}
		}else{
			w = x->p->lchild;
			if(w->color==RED){
				w->color = BLACK;
				x->p->color = RED;
				rightRotate(tree,x->p);
				w = x->p->lchild;
			}
			if(w->lchild->color==BLACK&&w->rchild->color==BLACK){
				w->color = RED;
				x = x->p;
			}else{
				if(w->lchild->color==BLACK){
					w->rchild->color = BLACK;
					w->color = RED;
					leftRotate(tree,w);
					w = x->p->lchild;
				}
				w->color = x->p->color;
				x->p->color = BLACK;
				w->lchild->color = BLACK;
				rightRotate(tree,x->p);
				x = tree->proot;
			}
		}
	}
	x->color = BLACK;
}

RBTreeNode* rbTreemin(RBTree *tree,RBTreeNode *z){
	RBTreeNode* x = z;
	while(z!=tree->pnil){
		if(z->lchild!=tree->pnil){
			x = z->lchild;
		}
		z=z->lchild;
	}
	return x;
}
void deleteNode(RBTree *tree,RBTreeNode *z){
	RBTreeNode *y = z;
	RBTreeNode *x;
	enum COLOR orginColor = y->color;
	
	if(z->lchild==tree->pnil){
		x = z->rchild;
		rbTransplant(tree,z,z->rchild);
	}else if(z->rchild==tree->pnil){
		x = z->lchild;
		rbTransplant(tree,z,z->lchild);
	}else{
		y = rbTreemin(tree,z->rchild);
		orginColor = y->color;
		x = y->rchild;
		if(y->p==z){
			x->p = y;
		}else{
			rbTransplant(tree,y,y->rchild);
			y->rchild = z->rchild;
			y->rchild->p = y;
		}
		rbTransplant(tree,z,y);
		y->lchild = z->lchild;
		y->lchild->p= y;
		y->color = z->color;
	}
	
	if(orginColor==BLACK){
		rbDeleteFixup(tree,x);
	}
	free(z);
}

int deleteRBTree(RBTree *tree,int key){
	
	RBTreeNode *x = tree->proot;
	while(x!=tree->pnil){
		if(x->key>key){
			x = x->lchild;
		}else if(x->key<key){
			x = x->rchild;
		}else{
			deleteNode(tree,x);
			return OK;
		}
	}
	return ERR;
}




void test(RBTree T){
	T.proot->color = RED;
}
int main(){
	srand((unsigned)time(NULL));
	int key=0;
	int i=0;
	int size = 10;
	int range = 100;

	printf("insert test==\n");
	//search test
	int arr[10];
	int j=0;
	int flag = 0;
	for(i=0;i<size;i++){
		key = rand()%range;
		flag = 0;
		for(j=0;j<i;j++){
			if(arr[j]==key){
				flag = 1;
				break;
			}
		}
		if(flag){
			i--;
		}else{
			arr[i]=key;
		}
	}

	RBTree tree;
	RBTreeNode nil;
	nil.key = 0;
	nil.color = BLACK;
	nil.lchild = NULL;
	nil.rchild = NULL;
	nil.p = NULL;
	tree.pnil = &nil;
	tree.proot = &nil;
	//test(tree);
	//printf("test %d\n",tree.proot->color);
	for(i=0;i<size;i++){
		insertRBTree(&tree,arr[i]);
    }
	
	
	printf("arr: [");
	for(i=0;i<size;i++){
		printf("%d ",arr[i]);
	}
	printf("]\n");
	tracePrint(tree.proot,tree.pnil,size);
	printf("\n");
	
	printf("delelte test==\ninput a index for delete begin from 1:\n");
	int index = -1;
	scanf("%d",&index);
	//delete test
	key = arr[index-1];
	if(!deleteRBTree(&tree,key)){
		printf("delete success\n");
	}
	key = arr[index-1]-1;
	if(!deleteRBTree(&tree,key)){
		printf("delete success\n");
	}else{
		printf("delete not found key: %d\n",key);
	}
	
	tracePrint(tree.proot,tree.pnil,size);
	printf("\n");
	
	return 0;
}